# -*- coding: utf-8 -*-
# @Time     : 2020/9/18 2:23 下午
# @Author   : huangxiong
# @FileName : regression_tools.py
# @Comment  : 风格分析的回归函数
# @Software : PyCharm
import os
os.environ['OMP_NUM_THREADS'] = '1'
import pandas as pd
import numpy as np

from sklearn.linear_model import Lasso, LassoCV
from sklearn.linear_model import Ridge, RidgeCV
from sklearn.metrics import r2_score
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.project_tool.math_func import math_tool
from quant_researcher.quant.project_tool import hammer
from quant_researcher.quant.project_tool.exception import FOFUserError, ERR_U_2


def get_regression_res_by_nav(start_end=(None, None), **kwargs):
    """
    输入基金、基金经理的收益率序列(日频、月频均可)，以及指数的收益率序列，按照指定的回归方式，得到回归结果

    :param tuple start_end: 回归起始和结束日期，用来对DataFrame进行索引
                            reg_start: 回归起始日期，默认为None，表示从第一行开始
                            reg_end: 回归结束日期，默认为None，表示回归到最后一行
                            若对整个DataFrame进行回归，请传(None, None)
                            由于滚动回归需要使用多进程，该参数必须为第一个参数，
                            而且希望单独使用该函数时不需要用户传该参数，所以该参数写成了关键字参数，其他参数都放在了**kwargs里
    :param kwargs:
        - pd.DataFrame df_to_fit: 回归的自变量和因变量，格式如下，
                                  列trade_date为日期，列ret为基金或基金经理收益率，其他为指数收益率
            +------------+---------+---------+---------+---------+---------+---------+
            | trade_date |   ret   |  000929 |  000930 |  000931 |  000936 |  000937 |
            +------------+---------+---------+---------+---------+---------+---------+
            | 2020-05-06 |  0.0017 |  0.0118 |  0.0115 |  0.0102 |  0.0288 | -0.0013 |
            | 2020-05-07 |  0.0073 |  0.0003 | -0.0038 | -0.0034 | -0.0061 | -0.0038 |
            | 2020-05-08 |  0.0159 |  0.0098 |  0.007  |  0.0246 |  0.0172 |  0.0026 |
            | 2020-05-11 |  0.0062 | -0.0024 |  0.0037 |  0.0058 |  0.0154 | -0.0005 |
            | 2020-05-12 |  0.0094 | -0.0052 |  -0.001 |  0.0054 | -0.0107 | -0.0035 |
            +------------+---------+---------+---------+---------+---------+---------+
        - str reg_method: 回归方法，目前支持minimize、lasso、ridge
        - bool need_intercept: 是否需要截距项
        - bool need_bounds: 是否需要给回归系数设置边界。minimize回归设置的边界为(0, 1)，lasso回归设置的边界为(0, ∞)
        - bool need_r_square: 是否需要R²
        - bool need_constraints: 是否需要设置约束条件，该参数只针对minimize回归，若需要，则系数之和为1。默认：True
        - float lasso_alpha: lasso回归时的alpha参数，若不传，则使用LassoCV寻找最适合的alpha
        - float ridge_alpha: ridge回归时的alpha参数，若不传，则使用RidgeCV寻找最适合的alpha
    :return: pd.DataFrame or None。如果回归不成功，返回None；如果回归成功，返回各指数的系数、截距项和R²
        +------------+--------------+------------+-----------+----------+
        | index_code | index_weight |  reg_date  | intercept | r_square |
        +------------+--------------+------------+-----------+----------+
        |   000929   |    0.3131    | 2020-06-16 |   0.0012  |  0.5283  |
        |   000930   |     0.0      | 2020-06-16 |   0.0012  |  0.5283  |
        |   000931   |    0.4053    | 2020-06-16 |   0.0012  |  0.5283  |
        |   000936   |     0.0      | 2020-06-16 |   0.0012  |  0.5283  |
        |   000937   |     0.0      | 2020-06-16 |   0.0012  |  0.5283  |
        +------------+--------------+------------+-----------+----------+
    """
    df_to_fit = kwargs.pop('df_to_fit')
    reg_method = kwargs.pop('reg_method')
    need_intercept = kwargs.pop('need_intercept', True)
    need_bounds = kwargs.pop('need_bounds', True)
    need_r_square = kwargs.pop('need_r_square', True)
    need_constraints = kwargs.pop('need_constraints', True)
    lasso_alpha = kwargs.pop('lasso_alpha', None)
    ridge_alpha = kwargs.pop('ridge_alpha', None)

    reg_start, reg_end = start_end
    intercept = r_square = np.nan
    if reg_start is not None:
        if reg_start >= df_to_fit['trade_date'].iloc[-1]:
            LOG.error(f"设置的回归起点为{reg_start}，DataFrame最大日期为{df_to_fit['trade_date'].iloc[-1]}"
                      f"前者应该比后者小，请重新输入参数")
            return
    df_to_fit = df_to_fit.ffill().bfill().set_index('trade_date')
    df_to_fit = df_to_fit.loc[reg_start: reg_end, :]
    if df_to_fit.shape[0] < 10:
        LOG.error(f"用来拟合的点小于10个，请检查输入参数，增加拟合的数据个数")
        return
    index_list = df_to_fit.columns[1:]
    x_mat = df_to_fit.iloc[:, 1:].values
    # 将reg_end作为回归日期
    if reg_method == 'minimize':
        y_mat = df_to_fit['ret'].values.reshape((-1, 1))

        minimize_res = math_tool.fit_with_constraints_by_array(
            x_s=x_mat, y_s=y_mat,
            with_intercept=need_intercept,
            no_bounds=not need_bounds,
            set_constraints=need_constraints,
            cpt_r_square=need_r_square,
        )
        if minimize_res is None:
            LOG.error(f'没有拟合成功！请检查数据是否正常！！！')
            return

        weights = minimize_res['weight'][0]
        res_df = pd.DataFrame({'index_code': index_list, 'index_weight': weights})
        if need_intercept:
            intercept = minimize_res['intercept']
        if need_r_square:
            r_square = minimize_res['r_square']
    elif reg_method == 'lasso':
        y_mat = df_to_fit['ret'].values.ravel()

        if lasso_alpha is None:
            lasso_cv = LassoCV(cv=5, positive=need_bounds, fit_intercept=need_intercept)
            lasso_cv.fit(x_mat, y_mat)
            lasso_alpha = lasso_cv.alpha_

        model = Lasso(alpha=lasso_alpha, positive=need_bounds, fit_intercept=need_intercept)
        model.fit(x_mat, y_mat)
        weights = model.coef_
        res_df = pd.DataFrame({'index_code': index_list, 'index_weight': weights})
        if need_intercept:
            intercept = model.intercept_
        if need_r_square:
            y_predict = model.predict(x_mat)
            r_square = r2_score(y_mat, y_predict)
    elif reg_method == 'ridge':
        y_mat = df_to_fit['ret'].values.ravel()

        if ridge_alpha is None:
            ridge_cv = RidgeCV(cv=5, fit_intercept=need_intercept)
            ridge_cv.fit(x_mat, y_mat)
            ridge_alpha = ridge_cv.alpha_

        model = Ridge(alpha=ridge_alpha, fit_intercept=need_intercept)
        model.fit(x_mat, y_mat)
        weights = model.coef_
        res_df = pd.DataFrame({'index_code': index_list, 'index_weight': weights})
        if need_intercept:
            intercept = model.intercept_
        if need_r_square:
            y_predict = model.predict(x_mat)
            r_square = r2_score(y_mat, y_predict)
    else:
        LOG.error(f"目前只支持minimize回归、lasso回归、ridge回归，请修改reg_method参数")
        raise NotImplementedError

    res_df['reg_date'] = reg_end
    if need_intercept:
        res_df['intercept'] = intercept
    if need_r_square:
        res_df['r_square'] = r_square

    return res_df


def get_rolling_regression_res_by_nav(x_df, y_df, reg_method, need_intercept, need_bounds,
                                      need_r_square, windows, **kwargs):
    """
    输入基金、基金经理的收益率序列(日频、月频均可)，以及指数的收益率序列，以及滚动回归窗口，按照指定的回归方式，得到滚动回归结果

    :param pd.DataFrame x_df: 回归的自变量，格式如下，列trade_date为日期，列ret为收益率
        +------------+---------+
        | trade_date |   ret   |
        +------------+---------+
        | 2020-05-07 |  0.0073 |
        | 2020-05-08 |  0.0159 |
        | 2020-05-11 |  0.0062 |
        | 2020-05-12 |  0.0094 |
        | 2020-05-13 |  0.0122 |
        +------------+---------+
    :param pd.DataFrame y_df: 回归的因变量，格式如下，列trade_date为日期(格式需与x_df的相同)，其他列名为指数代码，值为收益率
        +------------+---------+---------+---------+---------+---------+
        | trade_date |  000929 |  000930 |  000931 |  000936 |  000937 |
        +------------+---------+---------+---------+---------+---------+
        | 2020-05-07 |  0.0003 | -0.0038 | -0.0034 | -0.0061 | -0.0038 |
        | 2020-05-08 |  0.0098 |  0.007  |  0.0246 |  0.0172 |  0.0026 |
        | 2020-05-11 | -0.0024 |  0.0037 |  0.0058 |  0.0154 | -0.0005 |
        | 2020-05-12 | -0.0052 |  -0.001 |  0.0054 | -0.0108 | -0.0035 |
        | 2020-05-13 |  0.0036 |  0.0009 | -0.0017 | -0.0024 |  0.0024 |
        +------------+---------+---------+---------+---------+---------+
    :param str reg_method: 回归方法，目前支持minimize、lasso、ridge
    :param bool need_intercept: 是否需要截距项
    :param bool need_bounds: 是否需要给回归系数设置边界。minimize回归设置的边界为(0, 1)，lasso回归设置的边界为(0, ∞)
    :param bool need_r_square: 是否需要R²
    :param int windows: 滚动回归窗口。函数中是直接用windows的值对DataFrame进行切片。
                        若windows=10，输入的是日收益率序列，则表示滚动天数为10天；
                        若输入的是月收益率序列，则表示滚动天数为300天
    :param kwargs:
        - list reg_list，         多进程需要的切片参数，可以自己输入，如果不传，则用默认方法生成
        - bool need_constraints， 是否需要设置约束条件，该参数只针对minimize回归，若需要，则系数之和为1。默认：True
        - float lasso_alpha，     lasso回归时的alpha参数，若不传，则使用LassoCV寻找最适合的alpha
        - float ridge_alpha，     ridge回归时的alpha参数，若不传，则使用RidgeCV寻找最适合的alpha
    :return: pd.DataFrame。将滚动回归得到的各DataFrame合并
        +------------+--------------+------------+-----------+----------+
        | index_code | index_weight |  reg_date  | intercept | r_square |
        +------------+--------------+------------+-----------+----------+
        |   000929   |    0.3131    | 2020-06-16 |   0.0012  |  0.5283  |
        |   000930   |     0.0      | 2020-06-16 |   0.0012  |  0.5283  |
        |   000931   |    0.4053    | 2020-06-16 |   0.0012  |  0.5283  |
        |   000936   |     0.0      | 2020-06-16 |   0.0012  |  0.5283  |
        |   000937   |     0.0      | 2020-06-16 |   0.0012  |  0.5283  |
        +------------+--------------+------------+-----------+----------+
    """
    reg_list = kwargs.pop('reg_list', None)
    need_constraints = kwargs.pop('need_constraints', True)
    lasso_alpha = kwargs.pop('lasso_alpha', None)
    ridge_alpha = kwargs.pop('ridge_alpha', None)

    df_to_fit = y_df.merge(x_df, how='inner', on='trade_date').sort_values(by='trade_date')

    if windows > df_to_fit.shape[0]:
        raise FOFUserError(ERR_U_2, f"收益率序列的长度为{df_to_fit.shape[0]}，比滚动窗口的长度{windows}小，"
                                    f"无法进行滚动回归，建议缩小windows的值")

    if reg_list is None:
        date_list = df_to_fit['trade_date'].tolist()
        reg_end_list = date_list[windows:]
        reg_start_list = date_list[:df_to_fit.shape[0] - windows]
        reg_list = list(zip(reg_start_list, reg_end_list))

    reg_info = {
        'func': get_regression_res_by_nav,
        'df_to_fit': df_to_fit,
        'reg_method': reg_method,
        'need_intercept': need_intercept,
        'need_bounds': need_bounds,
        'need_r_square': need_r_square,
        'need_constraints': need_constraints,
        'lasso_alpha': lasso_alpha,
        'ridge_alpha': ridge_alpha,
    }

    result_list = hammer.split_a_list_into_batches_and_deal_each_with_a_process(
        the_list=reg_list, info=reg_info,
        ignore_err=True, how_many=8
    )

    if len(result_list) == 0:
        return
    else:
        result_df = pd.concat(result_list)
        return result_df


def test_func():
    from quant_researcher.quant.datasource_fetch.fund_api import fund_nav_related
    from quant_researcher.quant.datasource_fetch.index_api import index_price_related
    # f_df = fund_nav_related.get_fund_daily_return('110022', start='2020-05-01', end='2020-09-01',
    #                                               return_only_price_col=False)
    f_df = fund_nav_related.get_fund_return(
        fund_code='110022',
        start_date='2020-05-01',
        end_date='2020-09-01'
    )
    f_df['end_date'] = pd.to_datetime(f_df['end_date'])
    f_df['end_date'] = f_df['end_date'].dt.date
    f_df = f_df.rename(columns={"end_date": 'tj', "daily_return": 'f_price'})
    f_df['tj'] = f_df['tj'].astype(str)
    f_df = f_df[['tj', 'f_price']].rename(columns={'tj': 'trade_date', 'f_price': 'ret'})
    # i_df = index_price_related.get_index_daily_return(['000929', '000930', '000931', '000936', '000937'],
    #                                                   start='2020-05-01', end='2020-09-01',
    #                                                   return_only_price_col=False, replace_tj=True)
    i_df = index_price_related.get_index_return(
        index_code=['000929', '000930', '000931', '000936', '000937'],
        start_date='2020-05-01',
        end_date='2020-09-01'
    )
    i_df['end_date'] = pd.to_datetime(i_df['end_date'])
    i_df['end_date'] = i_df['end_date'].dt.date
    i_df = i_df.rename(columns={'end_date': 'trade_date', 'daily_return': 'i_price'})
    i_df['trade_date'] = i_df['trade_date'].astype(str)
    i_df = i_df.set_index(['trade_date', 'index_code'])['i_price'].unstack().reset_index()

    df_to_fit = f_df.merge(i_df, how='inner', on='trade_date').sort_values(by='trade_date')
    # reg_result = get_regression_res_by_nav(df_to_fit=df_to_fit, reg_method='lasso', need_intercept=True,
    #                                        need_bounds=True, need_constraints=True, need_r_square=True)
    reg_result = get_rolling_regression_res_by_nav(i_df, f_df, reg_method='lasso', need_intercept=True,
                                                   need_bounds=True, need_constraints=True, need_r_square=True,
                                                   windows=30)
    return reg_result


if __name__ == "__main__":
    test_func()
